 #
#   Movie of the probability density!
#

#
#   Copyright (C) 2012
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

import os
if not os.path.isdir("./data"):
    os.makedirs("./data")

def segments(min,max,steps):
    return [i/float(steps)*(max-min)+min for i in range(steps+1)]

from HarmonicOscillator1D import Psi as wvfn
from HarmonicOscillator1D import omega

from math import pi
timesteps=125
T_initial = 0.5*pi/omega
T_final = T_initial+2*(2*pi/omega) # A full `physical' period.

# Solve...
count = 0
outfile = open("./wvfn.csv",'w')

for time in segments(T_initial,T_final,timesteps): 
    print(count)
    count += 1
    int = 0
    for position in segments(-40,40,125):
        rho = abs(wvfn(time,position))**2
        int += rho*(40-(-40))/200.0
        outfile.write(str(time)+", "+str(position)+", "+str(rho)+", "+str(int)+"\n")
    
outfile.close()

print("Updated './data/'")